import pandas as pd
from sklearn.model_selection import StratifiedKFold, KFold
from sklearn.svm import SVR
from sklearn.model_selection import train_test_split
import numpy as np


class DataSets:
    def __init__(self, path_x,
                 path_y,
                 label_index=1,
                 label_threshold=0.014,
                 test_proportion=0.8):
        '''
            定义数据和标签的路径
            读取并处理数据格式
        :param path_x: 数据的路径
        :param path_y: 标签的路径
        '''
        self.path_x = path_x
        self.path_y = path_y

        self.test_proportion = test_proportion

        self.x = pd.read_csv(self.path_x, index_col=None, header=None).to_numpy()
        self.y_ori = pd.read_csv(self.path_y, index_col=None, header=None).to_numpy()[:, label_index]

        # 贴标签
        self.y = list(map(lambda x: 1 if x > label_threshold else 0, self.y_ori))
        self.y = np.asarray(self.y)

        self.x_shape = self.x.shape
        self.y_shape = self.y.shape

    def split_train_test(self):
        data_train, data_test, label_train, label_test = train_test_split(self.x,
                                                                          self.y,
                                                                          test_size=self.test_proportion,
                                                                          stratify=self.y)  # 保证相同的比例
        return data_train, data_test, label_train, label_test


if __name__ == "__main__":
    path_x = "data/data_naca0012/dv.csv"
    path_y = "data/data_naca0012/fc.csv"

    ds = DataSets(path_x=path_x,
                  path_y=path_y)
    ds.split_train_test()
